# Copyright 2022 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License

"""Tests for the library of QDWH-based singular value decomposition."""
import functools

import jax
import jax.numpy as jnp
import numpy as np
import scipy.linalg as osp_linalg
from absl.testing import absltest
from jax._src import test_util as jtu
from jax._src.lax import svd
from jax.config import config

config.parse_flags_with_absl()
_JAX_ENABLE_X64 = config.x64_enabled

# Input matrix data type for SvdTest.
_SVD_TEST_DTYPE = np.float64 if _JAX_ENABLE_X64 else np.float32

# Machine epsilon used by SvdTest.
_SVD_TEST_EPS = jnp.finfo(_SVD_TEST_DTYPE).eps

# SvdTest relative tolerance.
_SVD_RTOL = 1e-6 if _JAX_ENABLE_X64 else 1e-2

_MAX_LOG_CONDITION_NUM = 9 if _JAX_ENABLE_X64 else 4
from ..utils.timer_wrapper import jax_op_timer, partial_timed


@jtu.with_config(jax_numpy_rank_promotion="allow")
class SvdTest(jtu.JaxTestCase):
    @jtu.sample_product(
        [dict(m=m, n=n) for m, n in zip([2, 8, 10, 20], [4, 6, 10, 18])],
        log_cond=np.linspace(1, _MAX_LOG_CONDITION_NUM, 4),
        full_matrices=[True, False],
    )
    def testSvdWithRectangularInput(self, m, n, log_cond, full_matrices):
        """Tests SVD with rectangular input."""
        with jax.default_matmul_precision("float32"):
            a = np.random.uniform(low=0.3, high=0.9, size=(m, n)).astype(
                _SVD_TEST_DTYPE
            )
            u, s, v = osp_linalg.svd(a, full_matrices=False)
            cond = 10**log_cond
            s = jnp.linspace(cond, 1, min(m, n))
            a = (u * s) @ v
            a = a.astype(complex) * (1 + 1j)

            osp_linalg_fn = functools.partial(
                osp_linalg.svd, full_matrices=full_matrices
            )
            timer = jax_op_timer()
            with timer:
                actual_u, actual_s, actual_v = svd.svd(a, full_matrices=full_matrices)
                timer.gen.send((actual_u, actual_s, actual_v))

            k = min(m, n)
            if m > n:
                unitary_u = jnp.real(actual_u.T.conj() @ actual_u)
                unitary_v = jnp.real(actual_v.T.conj() @ actual_v)
                unitary_u_size = m if full_matrices else k
                unitary_v_size = k
            else:
                unitary_u = jnp.real(actual_u @ actual_u.T.conj())
                unitary_v = jnp.real(actual_v @ actual_v.T.conj())
                unitary_u_size = k
                unitary_v_size = n if full_matrices else k

            _, expected_s, _ = osp_linalg_fn(a)
        
            svd_fn = partial_timed(lambda a: svd.svd(a, full_matrices=full_matrices))
            args_maker = lambda: [a]

            with self.subTest("Test JIT compatibility"):
                self._CompileAndCheck(svd_fn, args_maker)

            with self.subTest("Test unitary u."):
                self.assertAllClose(
                    np.eye(unitary_u_size), unitary_u, rtol=_SVD_RTOL, atol=2e-3
                )

            with self.subTest("Test unitary v."):
                self.assertAllClose(
                    np.eye(unitary_v_size), unitary_v, rtol=_SVD_RTOL, atol=2e-3
                )

            with self.subTest("Test s."):
                self.assertAllClose(
                    expected_s, jnp.real(actual_s), rtol=_SVD_RTOL, atol=1e-6
                )

    @jtu.sample_product(
        [dict(m=m, n=n) for m, n in zip([50, 6], [3, 60])],
    )
    def testSvdWithSkinnyTallInput(self, m, n):
        """Tests SVD with skinny and tall input."""
        # Generates a skinny and tall input
        with jax.default_matmul_precision("float32"):
            np.random.seed(1235)
            a = np.random.randn(m, n).astype(_SVD_TEST_DTYPE)
            timer = jax_op_timer()
            with timer:
                u, s, v = svd.svd(a, full_matrices=False, hermitian=False)
                timer.gen.send((u, s, v))

            relative_diff = np.linalg.norm(a - (u * s) @ v) / np.linalg.norm(a)

            np.testing.assert_almost_equal(relative_diff, 1e-6, decimal=6)

    @jtu.sample_product(
        [dict(m=m, r=r) for m, r in zip([8, 8, 8, 10], [3, 5, 7, 9])],
        log_cond=np.linspace(1, 3, 3),
    )
    def testSvdWithOnRankDeficientInput(self, m, r, log_cond):
        """Tests SVD with rank-deficient input."""
        with jax.default_matmul_precision("float32"):
            a = jnp.triu(jnp.ones((m, m))).astype(_SVD_TEST_DTYPE)

            # Generates a rank-deficient input.
            timer = jax_op_timer()
            with timer:
                u, s, v = jnp.linalg.svd(a, full_matrices=False)
                timer.gen.send((u, s, v))
            cond = 10**log_cond
            s = jnp.linspace(cond, 1, m)
            s = s.at[r:m].set(0)
            a = (u * s) @ v

            with jax.default_matmul_precision("float32"):
                timer = jax_op_timer()
                with timer:
                    u, s, v = svd.svd(a, full_matrices=False, hermitian=False)
                    timer.gen.send((u, s, v))
            diff = np.linalg.norm(a - (u * s) @ v)

            np.testing.assert_almost_equal(diff, 1e-4, decimal=2)

    @jtu.sample_product(
        [dict(m=m, n=n) for m, n in zip([2, 4, 8], [4, 4, 6])],
        full_matrices=[True, False],
        compute_uv=[True, False],
        dtype=jtu.dtypes.floating + jtu.dtypes.complex,
    )
    def testSvdOnZero(self, m, n, full_matrices, compute_uv, dtype):
        """Tests SVD on matrix of all zeros."""
        osp_fun = functools.partial(
            osp_linalg.svd, full_matrices=full_matrices, compute_uv=compute_uv
        )
        lax_fun = partial_timed(
            svd.svd, full_matrices=full_matrices, compute_uv=compute_uv
        )
        args_maker_svd = lambda: [jnp.zeros((m, n), dtype=dtype)]
        self._CheckAgainstNumpy(osp_fun, lax_fun, args_maker_svd)
        self._CompileAndCheck(lax_fun, args_maker_svd)

    @jtu.sample_product(
        [
            dict(m=m, n=n, r=r, c=c)
            for m, n, r, c in zip([2, 4, 8], [4, 4, 6], [1, 0, 1], [1, 0, 1])
        ],
        dtype=jtu.dtypes.floating,
    )
    @jtu.skip_on_devices("rocm")
    def testSvdOnTinyElement(self, m, n, r, c, dtype):
        """Tests SVD on matrix of zeros and close-to-zero entries."""
        a = jnp.zeros((m, n), dtype=dtype)
        tiny_element = jnp.finfo(a).tiny
        a = a.at[r, c].set(tiny_element)

        @jax.jit
        def lax_fun(a):
            return svd.svd(a, full_matrices=False, compute_uv=False, hermitian=False)

        actual_s = lax_fun(a)

        k = min(m, n)
        expected_s = np.zeros((k,), dtype=dtype)
        expected_s[0] = tiny_element

        self.assertAllClose(expected_s, jnp.real(actual_s), rtol=_SVD_RTOL, atol=1e-6)


if __name__ == "__main__":
    absltest.main(testLoader=jtu.JaxTestLoader())
